import struct
import bitstring

import logging


# The default request size for blocks of pieces is 2^14 bytes.
#
# NOTE: The official specification states that 2^15 is the default request
#       size - but in reality all implementations use 2^14. See the
#       unofficial specification for more details on this matter.
#
#       https://wiki.theory.org/BitTorrentSpecification
#
REQUEST_SIZE = 2**14


class PeerMessage:
    """
    A message between two peers.

    All of the remaining messages in the protocol take the form of:
        <length prefix><message ID><payload>

    - The length prefix is a four byte big-endian value.
    - The message ID is a single decimal byte.
    - The payload is message dependent.

    NOTE: The Handshake messageis different in layout compared to the other
          messages.

    Read more:
        https://wiki.theory.org/BitTorrentSpecification#Messages

    BitTorrent uses Big-Endian (Network Byte Order) for all messages, this is
    declared as the first character being '>' in all pack / unpack calls to the
    Python's `struct` module.
    """
    Choke = 0
    Unchoke = 1
    Interested = 2
    NotInterested = 3
    Have = 4
    BitField = 5
    Request = 6
    Piece = 7
    Cancel = 8
    Port = 9
    Handshake = None  # Handshake is not really part of the messages
    KeepAlive = None  # Keep-alive has no ID according to spec

    def encode(self) -> bytes:
        """
        Encodes this object instance to the raw bytes representing the entire
        message (ready to be transmitted).
        """
        pass

    @classmethod
    def decode(cls, data: bytes):
        """
        Decodes the given BitTorrent message into a instance for the
        implementing type.
        """
        pass


class KeepAlive(PeerMessage):
    """
    The choke message is used to tell the other peer to stop send request
    messages until unchoked.

    Message format:
        <len=0001><id=0>
    """

    def encode(self) -> bytes:
        return struct.pack('>I', 
                           0,
                           PeerMessage.KeepAlive)
    
    @classmethod
    def decode(cls, data: bytes):
        return cls()

    def __str__(self) -> str:
        return "KeepAlive message"


class Choke(PeerMessage):
    """
    The choke message is used to tell the other peer to stop send request
    messages until unchoked.

    Message format:
        <len=0001><id=0>
    """

    def encode(self) -> bytes:
        return struct.pack('>Ib', 
                           1,
                           PeerMessage.Choke)
    
    @classmethod
    def decode(cls, data: bytes):
        return cls()

    def __str__(self) -> str:
        return "Chocke message"

class Unchoke(PeerMessage):
    """
    Unchoking a peer enables that peer to start requesting pieces from the
    remote peer.

    Message format:
        <len=0001><id=1>
    """

    def encode(self) -> bytes:
        return struct.pack('>Ib', 
                           1,
                           PeerMessage.Unchoke)
    
    @classmethod
    def decode(cls, data: bytes):
        return cls()
    
    def __str__(self) -> str:
        return "Unchocke message"

class Interested(PeerMessage):
    """
    The interested message is fix length and has no payload other than the
    message identifiers. It is used to notify each other about interest in
    downloading pieces.

    Message format:
        <len=0001><id=2>
    """
    def encode(self) -> bytes:
        return struct.pack('>Ib', 
                           1,
                           PeerMessage.Interested)
    
    @classmethod
    def decode(cls, data: bytes):
        return cls()
    
    def __str__(self) -> str:
        return 'Interested message'


class NotInterested(PeerMessage):
    """
    The interested message is fix length and has no payload other than the
    message identifiers. It is used to notify each other about interest in
    downloading pieces.

    Message format:
        <len=0001><id=2>
    """
    def encode(self) -> bytes:
        return struct.pack('>Ib', 
                           1,
                           PeerMessage.NotInterested)
    
    @classmethod
    def decode(cls, data: bytes):
        return cls
    
    def __str__(self) -> str:
        return 'NotInterested message'

class Have(PeerMessage):
    """
    Represents a piece successfully downloaded by the remote peer. The piece
    is a zero based index of the torrents pieces

    Message format:
        <len=0005><id=4><piece index>
    """
    def __init__(self, index: int):
        self.index = index

    def encode(self) -> bytes:
        return struct.pack('>IbI',
                           5,
                           PeerMessage.Have,
                           self.index)

    @classmethod
    def decode(cls, data: bytes):
        logging.debug('Decoding a Have message of length: {}'.format(len(data)))

        parts = struct.unpack('>IbI', data)
        index = parts[2]
        return cls(index)
    
    def __str__(self) -> str:
        return 'Have message [index: {}]'.format(self.index)


class BitField(PeerMessage):
    """
    The BitField is a message with variable length where the payload is a
    bit array representing all the bits a peer have (1) or does not have (0).

    Message format:
        <len=0001+X><id=5><bitfield>
    """
    def __init__(self, data):
        self.bitfield = bitstring.BitArray(bytes=data)
    
    @classmethod
    def decode(cls, data: bytes):
        logging.debug('Decoding a BitField message of length: {}'.format(len(data)))

        message_length = struct.unpack('>I', data[:4])[0]
        logging.debug('BitField message carry data with length: {}'.format(message_length))
        parts = struct.unpack('>Ib{}s'.format(message_length - 1), data)

        return cls(parts[2])

    def __str__(self) -> str:
        return 'BitField message with {} of file'.format(self.bitfield.count(1) / len(self.bitfield))
    
    def __str__(self) -> str:
        return 'BitField message [bitfield: {} / {}, {}]'.format(
            self.bitfield.count(1),
            len(self.bitfield),
            self.bitfield.count(1) / len(self.bitfield))


class Request(PeerMessage):
    """
    The message used to request a block of a piece (i.e. a partial piece).

    The request size for each block is 2^14 bytes, except the final block
    that might be smaller (since not all pieces might be evenly divided by the
    request size).

    Message format:
        <len=0013><id=6><index><begin><length>
    """
    def __init__(self, index: int, begin: int, length: int = REQUEST_SIZE):
        self.index = index
        self.begin = begin
        self.length = length
    
    def encode(self) -> bytes:
        return struct.pack('>IbIII',
                           13,
                           PeerMessage.Request,
                           self.index,
                           self.begin,
                           self.length)

    @classmethod
    def decode(cls, data: bytes):
        logging.debug('Decoding a Request message of length: {}'.format(len(data)))
        parts = struct.unpack('>IbIII', data)

        return cls(parts[2], parts[3], parts[4]) 
    
    def __str__(self) -> str:
        return 'Request message [index: {}, beigin: {}, length:{}]'.format(self.index, self.begin, self.length)


class Piece(PeerMessage):
    """
    The message carray piece

    Message format:
        <len=0009+X><id=7><index><begin><block>
    """
    def __init__(self, index: int, begin: int, block: bytes):
        self.index = index
        self.begin = begin
        self.block = block
    
    def encode(self) -> bytes:
        message_length = len(self.block)

        return struct.pack('>IbII{}s'.format(message_length),
                           9 + message_length,
                           self.index,
                           self.begin,
                           self.block)
    
    @classmethod
    def decode(cls, data: bytes):
        logging.debug('Decoding a Piece message of length: {}'.format(len(data)))

        message_length = struct.unpack('>I', data[:4])[0]
        logging.debug('Piece message carry data with length: {}'.format(message_length))
        parts = struct.unpack('>IbII{}s'.format(message_length -1 -4 -4 ), data)

        return cls(parts[2], parts[3], parts[4]) 
    
    def __str__(self) -> str:
        return 'Piece message [index: {}, beigin: {}, block({})]'.format(self.index, self.begin, len(self.block))


class Cancel(PeerMessage):
    """
    The cancel message.

    Message format:
        <len=0013><id=8><index><begin><length>
    """
    def __init__(self, index: int, begin: int, length: int = REQUEST_SIZE):
        self.index = index
        self.begin = begin
        self.length = length
    
    def encode(self) -> bytes:
        return struct.pack('>IbIII',
                           13,
                           PeerMessage.Cancel,
                           self.index,
                           self.begin,
                           self.length)

    @classmethod
    def decode(cls, data: bytes):
        logging.debug('Decoding a Cancel message of length: {}'.format(len(data)))
        parts = struct.unpack('>IbIII', data)

        return cls(parts[2], parts[3], parts[4]) 
    
    def __str__(self) -> str:
        return 'Cancel message [index: {}, beigin: {}, length: {}]'.format(self.index, self.begin, self.length)

class Handshake(PeerMessage):
    """
    The handshake message is the first message sent and then received from a
    remote peer.

    The messages is always 68 bytes long (for this version of BitTorrent
    protocol).

    Message format:
        <pstrlen><pstr><reserved><info_hash><peer_id>

    In version 1.0 of the BitTorrent protocol:
        pstrlen = 19
        pstr = "BitTorrent protocol".

    Thus length is:
        49 + len(pstr) = 68 bytes long.
    """
    HANDSHAKE_LENGTH = 19 + 49

    def __init__(self, info_hash: bytes, peer_id: bytes) -> None:
        super().__init__()
        assert type(info_hash) == bytes
        assert type(peer_id) == bytes
        self.info_hash = info_hash
        self.peer_id = peer_id

    def encode(self) -> bytes:
        return struct.pack(
            '>B19s8x20s20s',
            19,                         # Single byte (B)
            b'BitTorrent protocol',     # String 19s
                                        # Reserved 8x (pad byte, no value)
            self.info_hash,             # String 20s
            self.peer_id                # String 20s
        )
    
    @classmethod
    def decode(cls, data: bytes):
        logging.debug('Decoding a Handshake message of length: {}'.format(len(data)))

        if len(data) < Handshake.HANDSHAKE_LENGTH:
            return None
        parts = struct.unpack('>B19s8x20s20s', data)
        return cls(info_hash=parts[2], peer_id=parts[3])
    
    def __str__(self) -> str:
        return 'Handshake message [peer_id: {}, info_hash: {}]'.format(self.peer_id, self.info_hash)


